import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os

# 定义实际的曲线
def actual_curve(t):
    x = np.sin(np.cos(t)) * np.cos(np.sin(t))
    y = np.sin(np.cos(t)) * np.sin(np.sin(t))
    z = np.cos(np.cos(t))
    return x, y, z

# 定义一个函数来加载拟合数据
def load_data(filename):
    data = np.loadtxt(filename, delimiter=',')
    return data[:, 0], data[:, 1], data[:, 2]

# 画图函数
def plot_combined_curves(filenames, title):
    # 生成实际曲线数据
    t_values = np.linspace(0, 2 * np.pi, 500)
    actual_x, actual_y, actual_z = actual_curve(t_values)

    # 创建 3D 图形
    fig = plt.figure(figsize=(10, 7))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制实际曲线
    ax.plot(actual_x, actual_y, actual_z, label="Actual Curve", color='black', linestyle='dashed')

    # 颜色列表
    colors = ['red', 'green', 'blue', 'orange', 'purple', 'brown', 'pink', 'gray', 'cyan', 'magenta']

    # 绘制拟合曲线
    for i, filename in enumerate(filenames):
        x, y, z = load_data(filename)
        n = int(filename.split('N')[1].split('.')[0])
        color = colors[i % len(colors)]  # 循环使用颜色列表中的颜色
        ax.plot(x, y, z, label=f"Fitted Curve N={n}", color=color, linestyle='solid')
    
    # 设置标签和标题
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title(title)
    
    # 显示图例
    ax.legend()
    
    # 保存图形到 ../figure 目录
    output_dir = "../figure"
    os.makedirs(output_dir, exist_ok=True)  # 确保目录存在
    output_file = os.path.join(output_dir, f"{title}.png")
    plt.savefig(output_file, dpi=300)  # 保存为 PNG 文件，dpi=300 是较高的分辨率
    print(f"Saved plot as {output_file}")
    
    # 显示图形
    plt.show()

# 主函数
def main():
    # 文件名列表，根据实际生成的文件调整
    equidistant_files = [
        "E3_fitted_equidistant_N10.txt",
        "E3_fitted_equidistant_N40.txt",
        "E3_fitted_equidistant_N160.txt"
    ]
    
    chordal_files = [
        "E3_fitted_chordal_N10.txt",
        "E3_fitted_chordal_N40.txt",
        "E3_fitted_chordal_N160.txt"
    ]
    
    # 对每组文件画图
    plot_combined_curves(equidistant_files, "Equidistant Fitted Curves")
    plot_combined_curves(chordal_files, "Chordal Fitted Curves")

if __name__ == "__main__":
    main()